#!/usr/bin/env python3

import cv2 as cv
import numpy as np
from matplotlib import pyplot as plt

img = cv.imread('../pic/7.png', 0)
# opencv中的傅里叶变换
dft = cv.dft(np.float32(img), flags=cv.DFT_COMPLEX_OUTPUT)
dft_shift = np.fft.fftshift(dft)
# cv.magnitude函数传入实部和虚部，计算幅度谱
magnitude_spectrum = 20 * np.log(cv.magnitude(dft_shift[:, :, 0], dft_shift[:, :, 1]))

plt.subplot(121), plt.imshow(img, cmap='gray')
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
plt.subplot(122), plt.imshow(magnitude_spectrum, cmap='gray')
plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])
plt.show(block=False)
plt.waitforbuttonpress()
plt.close()

# 逆傅里叶变换
rows, cols = img.shape
# 只能接受整数
crow, ccol = rows // 2, cols // 2
# 首先创建一个掩码，中心正方形为1，其余全为零   即低频可通过，高频无法通过
mask = np.zeros((rows, cols, 2), np.uint8)
mask[crow - 30:crow + 30, ccol - 30:ccol + 30] = 1
# 并非矩阵乘法，而是Hadamard积，仅对应位置相乘   即原mask中0的位置全部使得fshift变为0，原mask中1的位置使得fshift变成dft_shif的值
fshift = dft_shift * mask
# fshift中仅剩低频信息
f_ishift = np.fft.ifftshift(fshift)
img_back = cv.idft(f_ishift)
img_back = cv.magnitude(img_back[:, :, 0], img_back[:, :, 1])

plt.subplot(121), plt.imshow(img, cmap='gray')
plt.title('Input Image'), plt.xticks([]), plt.yticks([])
plt.subplot(122), plt.imshow(img_back, cmap='gray')
plt.title('Magnitude Spectrum'), plt.xticks([]), plt.yticks([])

plt.show(block=False)
plt.waitforbuttonpress()
plt.close()
